{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  import pydrake\n",
    "  import underactuated\n",
    "except ImportError:\n",
    "  !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "  from jupyter_setup import setup_underactuated\n",
    "  setup_underactuated()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Rimless Wheel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from IPython import get_ipython\n",
    "\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "from pydrake.all import (DiagramBuilder, PlanarSceneGraphVisualizer, SceneGraph,\n",
    "                         Simulator)\n",
    "from pydrake.examples.rimless_wheel import (RimlessWheel, RimlessWheelGeometry,\n",
    "                                            RimlessWheelParams)\n",
    "\n",
    "def simulate(slope=0.08, initial_angle=0, initial_angular_velocity=5.0):\n",
    "  params = RimlessWheelParams()\n",
    "  params.set_slope(slope)\n",
    "\n",
    "  builder = DiagramBuilder()\n",
    "  rimless_wheel = builder.AddSystem(RimlessWheel())\n",
    "  scene_graph = builder.AddSystem(SceneGraph())\n",
    "  RimlessWheelGeometry.AddToBuilder(\n",
    "      builder, rimless_wheel.get_floating_base_state_output_port(), params,\n",
    "      scene_graph)\n",
    "  visualizer = builder.AddSystem(\n",
    "      PlanarSceneGraphVisualizer(scene_graph, xlim=[-8., 8.], ylim=[-2., 3.],\n",
    "                                 show=plt_is_interactive))\n",
    "  builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
    "                  visualizer.get_input_port(0))\n",
    "\n",
    "  diagram = builder.Build()\n",
    "  simulator = Simulator(diagram)\n",
    "\n",
    "  context = simulator.get_mutable_context()\n",
    "\n",
    "  diagram.GetMutableSubsystemContext(\n",
    "      rimless_wheel, context).get_numeric_parameter(0).set_slope(slope)\n",
    "  context.SetAccuracy(1e-4)\n",
    "  context.SetContinuousState([initial_angle, initial_angular_velocity])\n",
    "  simulator.Initialize()\n",
    "  \n",
    "  duration = 5.0 if get_ipython() else 0.1 # sets a shorter duration during testing\n",
    "  AdvanceToAndVisualize(simulator, visualizer, duration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulate(initial_angular_velocity=5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulate(initial_angular_velocity=10.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulate(initial_angular_velocity=0.95)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulate(initial_angular_velocity=-5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulate(initial_angular_velocity=-4.8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a little interactive plot to allow you to visualize the trajectories of the rimless wheel as you vary the initial velocity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from ipywidgets import interact, widgets\n",
    "\n",
    "from underactuated.jupyter import SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "from pydrake.all import Simulator\n",
    "from pydrake.examples.rimless_wheel import RimlessWheel\n",
    "\n",
    "rimless_wheel = RimlessWheel()\n",
    "simulator = Simulator(rimless_wheel)\n",
    "context = simulator.get_mutable_context()\n",
    "params = context.get_numeric_parameter(0)\n",
    "qmin = params.slope() - rimless_wheel.calc_alpha(params) - 0.1\n",
    "qmax = params.slope() + rimless_wheel.calc_alpha(params) + 0.1\n",
    "qdmin = -2\n",
    "qdmax = 2\n",
    "\n",
    "context.SetAccuracy(1e-2)\n",
    "integrator = simulator.get_mutable_integrator()\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "# TODO(russt): make the slope interactive, too.\n",
    "def simulate(initial_angular_velocity=1.5, duration=1.5):\n",
    "    rimless_wheel.SetDefaultContext(context)\n",
    "    context.SetTime(0.0)\n",
    "    if initial_angular_velocity >= 0:\n",
    "        initial_angle = params.slope() - rimless_wheel.calc_alpha(params)\n",
    "    else:\n",
    "        initial_angle = params.slope() + rimless_wheel.calc_alpha(params)\n",
    "    if initial_angular_velocity == 0:\n",
    "        # Set double_support = True.\n",
    "        context.get_mutable_abstract_state(0).set_value(True)\n",
    "        \n",
    "    context.SetContinuousState([initial_angle, initial_angular_velocity])\n",
    "\n",
    "    integrator.StartDenseIntegration()\n",
    "    simulator.AdvanceTo(duration)\n",
    "    pp = integrator.StopDenseIntegration()\n",
    "\n",
    "    return pp.vector_values(pp.get_segment_times())\n",
    "\n",
    "if plt_is_interactive:\n",
    "    data = simulate()\n",
    "    line, = ax.plot(data[0,:], data[1,:],'b')\n",
    "    pt, = ax.plot(data[0,0], data[1,0],'b*', markersize=12)\n",
    "    \n",
    "    def update(initial_angular_velocity):\n",
    "        data = simulate(initial_angular_velocity)\n",
    "        line.set_xdata(data[0,:])\n",
    "        line.set_ydata(data[1,:])\n",
    "        pt.set_xdata(data[0,0])\n",
    "        pt.set_ydata(data[1,0])\n",
    "        fig.canvas.draw()\n",
    "    \n",
    "    interact(update, initial_angular_velocity=widgets.FloatSlider(min=qdmin, max=qdmax, step=.1, value=1.1))\n",
    "\n",
    "else:\n",
    "    data = simulate()\n",
    "    ax.plot(data[0,:], data[1,:],'b')\n",
    "    ax.plot(data[0,0], data[1,0],'b*', markersize=12)\n",
    "\n",
    "# Plot the energy contours.\n",
    "nq = 151\n",
    "nqd = 151\n",
    "mgl = params.mass() * params.gravity() * params.length()\n",
    "q = np.linspace(qmin, qmax, nq)\n",
    "qd = np.linspace(qdmin, qdmax, nqd)\n",
    "Q, QD = np.meshgrid(q, qd)\n",
    "Energy = .5 * params.mass() * params.length()**2 * QD**2 + mgl * np.cos(Q)\n",
    "ax.contour(Q,\n",
    "           QD,\n",
    "           Energy,\n",
    "           alpha=0.5,\n",
    "           linestyles=\"dashed\",\n",
    "           colors=\"black\",\n",
    "           linewidths=0.5)\n",
    "\n",
    "ax.plot(params.slope() - rimless_wheel.calc_alpha(params)*np.array([1, 1]), np.array([0, qdmax]), 'k--')\n",
    "ax.plot(params.slope() - rimless_wheel.calc_alpha(params)*np.array([1, 1]), np.array([0, qdmin]), 'k', linewidth=.25)\n",
    "ax.plot(params.slope() + rimless_wheel.calc_alpha(params)*np.array([1, 1]), np.array([0, qdmin]), 'k--')\n",
    "ax.plot(params.slope() + rimless_wheel.calc_alpha(params)*np.array([1, 1]), np.array([0, qdmax]), 'k', linewidth=.25)\n",
    "ax.plot([qmin, qmax], [0, 0], 'k', linewidth=.25)\n",
    "ax.plot([0, 0], [qdmin, qdmax], 'k', linewidth=.25)\n",
    "ax.set_xlabel(\"theta\")\n",
    "ax.set_ylabel(\"thetadot\")\n",
    "ax.axis([qmin, qmax, qdmin, qdmax])\n",
    "ax.set_title(\"Trajectories of the Rimless Wheel (w/ contours of \"\n",
    "             \"constant energy)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Compass Gait"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython import get_ipython\n",
    "\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "from pydrake.all import (ConstantVectorSource, DiagramBuilder,\n",
    "                         PlanarSceneGraphVisualizer, SceneGraph, SignalLogger,\n",
    "                         Simulator)\n",
    "from pydrake.examples.compass_gait import (CompassGait, CompassGaitGeometry,\n",
    "                                           CompassGaitParams)\n",
    "\n",
    "builder = DiagramBuilder()\n",
    "compass_gait = builder.AddSystem(CompassGait())\n",
    "\n",
    "hip_torque = builder.AddSystem(ConstantVectorSource([0.0]))\n",
    "builder.Connect(hip_torque.get_output_port(0), compass_gait.get_input_port(0))\n",
    "\n",
    "scene_graph = builder.AddSystem(SceneGraph())\n",
    "CompassGaitGeometry.AddToBuilder(\n",
    "    builder, compass_gait.get_floating_base_state_output_port(), scene_graph)\n",
    "visualizer = builder.AddSystem(\n",
    "    PlanarSceneGraphVisualizer(scene_graph, xlim=[-1., 5.], ylim=[-1., 2.], \n",
    "                               show=plt_is_interactive))\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
    "                visualizer.get_input_port(0))\n",
    "\n",
    "logger = builder.AddSystem(SignalLogger(14))\n",
    "builder.Connect(compass_gait.get_output_port(1), logger.get_input_port(0))\n",
    "\n",
    "diagram = builder.Build()\n",
    "simulator = Simulator(diagram)\n",
    "\n",
    "context = simulator.get_mutable_context()\n",
    "context.SetAccuracy(1e-4)\n",
    "context.SetContinuousState([0., 0., 0.4, -2.])\n",
    "\n",
    "duration = 8.0 if get_ipython() else 0.1 # sets a shorter duration during testing\n",
    "AdvanceToAndVisualize(simulator, visualizer, duration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(logger.data()[4, :], logger.data()[11, :])\n",
    "plt.xlabel(\"left leg angle\")\n",
    "plt.ylabel(\"left leg angular velocity\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
